# find optimal spectrum on arbitrary grid for a real function
from numpy import *
from math import pi
#import scipy.linalg.lstsq as LSQR
import matplotlib.pyplot as plt
from cplinalg import *
from cpsignalproc import *
def optimspecr(t,x,w):
	N=t.size
	L=w.size
	dt=(t[-1]-t[0])/(t.size-1.0)
	dw=(w[-1]-w[0])/(w.size-1.0)
	exl=ceil(2*pi/dw/dt)
	M=zeros((N,2*L))
        cf=zeros(2*L)
	for i in range(0,L):
		M[:,i]=sin(w[i]*t).T
		M[:,i+L]=cos(w[i]*t).T
#	sv=linalg.svd(M[:,1:])[1]
#	print sv
#	print sv.size, M.shape[1]
#	print rank(mat(M[:,1:]))
	avamp=sqrt(sum(x**2)/L)
	err=1
	nerr=1
	conv=1
	rx=x
	pfi=argmax(abs(fft.rfft(rx)))
	pf=2*pi/N/dt*pfi
	if pf>w[-1]:
		rx=rechighcut(rx,w[-1]+dw,dt,exl)[0:N]
	if pf<w[0]:
		rx=recbndpass(rx,w[0]-dw,w[-1]+dw,dt,exl)[0:N]
	mrx=max(abs(rx))
	y=zeros(w.size)*1j
	if avamp==0:
		return y
#	plt.plot(t,rx)
#	plt.show()
	rnd=0
	while err>1e-8 and mrx>10000*finfo(float).eps:
		if w[0]==0:
#			cf[1:]=linalg.lstsq(M[:,1:],rx.T,rcond=1e-12)[0]
#			cf[1:]=LSQR(M[:,1:],rx.T)[0]
			cf[1:]=tiklstsq(M[:,1:],rx.T,1e-8)
			cf[0]=0
		else:
#			cf=linalg.lstsq(M,rx.T,rcond=1e-12)[0]
#			cf=LSQR(M,rx.T)[0]
			cf=tiklstsq(M,rx.T,1e-8)
		xf=cf[L:2*L]-1j*cf[0:L]
		imax=argmax(abs(xf))
		if w[imax]!=0:
			fitv=real(xf[imax])*cos(w[imax]*t)-imag(xf[imax])*sin(w[imax]*t)
			y[imax]=y[imax]+xf[imax]
		else:
			fitv=real(xf[imax])
			y[imax]=y[imax]+real(xf[imax])
		rx=rx-fitv
		mrx=max(abs(rx))
		nerr=mrx/avamp
		conv=abs(err-nerr)/avamp
		err=nerr		
		pfi=argmax(abs(fft.rfft(rx)))
		pf=2*pi/N/dt*pfi
		if pf>w[-1]:
			rx=rechighcut(rx,w[-1]+dw,dt,exl)[0:N]
		if pf<w[0]:
			rx=recbndpass(rx,w[0]-dw,w[-1]+dw,dt,exl)[0:N]
#		if rnd>=0 and rnd<1:
#			print "w[imax] is %f" % (w[imax])
#			print t,w
#			plt.plot(w,abs(xf))
#		if rnd==1:
#			plt.show()
		print "error is %e" % (err)
		print "max residual is %e" % (mrx)
#		print "w[imax] is %f" % (w[imax])
#		print "conv is %f" % (conv)
		rnd=rnd+1
#	plt.show()
	return y

def optimspecr2(t,x,w):
	if t.size!=x.size:
		print "The sizes of t and x don't match!"
		return False
	N=t.size
	L=w.size
	dt=(t[-1]-t[0])/(t.size-1.0)
	dw=(w[-1]-w[0])/(w.size-1.0)
	exl=ceil(2*pi/dw/dt)
	M=zeros((N,2,L))
	for i in range(0,L):
		M[:,:,i]=hstack([sin(w[i]*t).reshape(N,1),cos(w[i]*t).reshape(N,1)])
	ores=sum(x**2)	
	avamp=sqrt(ores/float(N))
	err=1.0
	conv=1.0
	rx=x
	pfi=argmax(abs(fft.rfft(rx)))
	pf=2*pi/N*pfi
	if pf>w[-1]:
		rx=rechighcut(rx,w[-1]+dw,dt,exl)[0:N]
	if pf<w[0]:
		rx=recbndpass(rx,w[0]-dw,w[-1]+dw,dt,exl)[0:N]
	mrx=max(abs(rx))
	y=zeros(w.size)*1j
	cf=array([[0.0],[0.0]])
#	plt.plot(t,rx)	
	rnd=0
	if avamp==0:
		return y
	while err>1e-6 and mrx>10000*finfo(float).eps and conv>1e-12:
		im=0
		mxf=0.0+1j*0.0
		mres=0
		for i in range(0,L):
			if w[i]!=0:
			        cf,res=linalg.lstsq(M[:,:,i],rx.reshape(N,1),rcond=1e-12)[0:2]
			else:
				cf[0]=0
				cf[1]=sum(rx)/float(N)
				res=sum((rx-cf[1])**2)
			cfrshp=cf.reshape(2,)
			if i==0:
				mres=res
				im=i
				mxf=cfrshp[1]-1j*cfrshp[0]
			else:
				if res<mres:
					im=i
					mres=res
					mxf=cfrshp[1]-1j*cfrshp[0]
		if w[im]!=0:
			fitv=dot(M[:,:,im],array([[-imag(mxf)],[real(mxf)]])).reshape(N,)
			y[im]=y[im]+mxf
		else:
			fitv=real(mxf)
			y[im]=y[im]+real(mxf)
		rx=rx-fitv 
		mrx=max(abs(rx))
		err=mrx/avamp
		conv=abs((ores)**0.5-(mres)**0.5)/avamp/N
		ores=mres
		print "exl is:", exl
		print "the im is:", im
		print "the mres is:", mres
		print "the mrx is:",mrx
		print "the conv is:",conv
		print "the err is:", err
#		if rnd<10:
#			plt.plot(t,rx)
#		if rnd==10:
#			plt.show()	
		rnd=rnd+1			
	return y

def optimspecr3(t,x,w):
	N=t.size
	L=w.size
	M=zeros((N,2*L))
        cf=zeros(2*L)
	for i in range(0,L):
		M[:,i]=sin(w[i]*t).T
		M[:,i+L]=cos(w[i]*t).T

	avamp=sqrt(sum(x**2)/N)
	err=1
	nerr=1
	conv=1
	rx=x
	mrx=max(abs(rx))
	y=zeros(w.size)*1j
	if avamp==0:
		return y
#	plt.plot(w,abs(xf))
#	plt.show()
	rnd=0
	while err>1e-6 and mrx>10000*finfo(float).eps:			
		if w[0]==0:
#			cf[1:]=linalg.lstsq(M[:,1:],rx.T,rcond=1e-8)[0]
#			cf[1:]=LSQR(M[:,1:],rx.T)[0]
			cf[1:]=tiklstsq(M[:,1:],rx.T)
			cf[0]=0
		else:
#			cf=linalg.lstsq(M,rx.T,rcond=1e-8)[0]
#			cf=LSQR(M,rx.T)[0]
			cf=tiklstsq(M,rx.T)
		xf=cf[L:2*L]-1j*cf[0:L]
		y=y+xf
		cy=vstack([-imag(y).reshape([L,1]),real(y).reshape([L,1])])
		rx=x-dot(M,cy).reshape(N)
		mrx=max(abs(rx))
		nerr=mrx/avamp
		conv=abs(err-nerr)/avamp
		err=nerr	
		rnd=rnd+1
#		if rnd>0 and rnd<10:
#			print "w[imax] is %f" % (w[imax])
#			plt.plot(w,abs(y))
#		if rnd==10:
#			plt.show()
#		print "error is %e" % (err)
#		print "max residual is %e" % (mrx)
#		print "w[imax] is %f" % (w[imax])
#		print "conv is %f" % (conv)
#	plt.show()
	return y

def optimspecr4(t,x,w):
	N=t.size
	dt=(t[-1]-t[0])/(t.size-1.0)
	dw=(w[-1]-w[0])/(w.size-1.0)
	tt=t-t[0]
	avamp=sqrt(sum(x**2)/N)
	err=1
	nerr=1
	conv=1
	rx=x.copy()
#	plt.plot(rx)
	mrx=max(abs(rx.reshape(rx.size,)))
	exl=int(max([ceil(2*pi/dw/dt),N]))
	y=zeros(w.size)*1j
	rnd=0
	if avamp==0:
		return y
	
	while err>1e-12 and mrx>10000*finfo(float).eps and conv>1e-16:
		rxf=fft.rfft(rx,exl)
		mi=argmax(abs(rxf))
		mxw=mi*2*pi/exl
		if mxw<w[0] or mxw>w[-1]:
			if w[0]>0:
				rx=recbndpass(rx,w[0]-dw,w[-1]+dw,dt,exl)
			else:
				rx=rechighcut(rx,w[-1]+dw,dt,exl)
			mrx=max(abs(rx.reshape(rx.size,)))
			nerr=mrx/avamp
			conv=1
			err=nerr
		else:
#			if mxw==0:
#				M=cos(mxw*t).reshape(N,1)
#			else:
#				M=hstack([sin(mxw*t).reshape(N,1),cos(mxw*t).reshape(N,1)])
#			cf,res=linalg.lstsq(M,rx[0:N].reshape(N,1),rcond=1e-12)[0:2]
#			cf=cf.reshape(cf.size,)
#			if cf.size==1:
#				fitv=cf[0]
#				rx=rx[0:N]-fitv
#				y[mi]=y[mi]+cf[0]
#			else:
#				fitv=cf[0]*sin(mxw*t)+cf[1]*cos(mxw*t)
#				rx=rx[0:N]-fitv
#				y[mi]=y[mi]+cf[1]-1j*cf[0]
			fitv=real(rxf[mi])*2/exl*cos(mxw*tt)-imag(rxf[mi])*2/exl*sin(mxw*tt)
			rx=rx[0:N]-fitv
			y[mi]=y[mi]+rxf[mi]
			mrx=max(abs(rx.reshape(rx.size,)))
			nerr=mrx/avamp
#			print "err is:", err
#			print "nerr is:", nerr
			conv=abs(err-nerr)/avamp
			err=nerr
#		if rnd<20:
#			plt.plot(rx)
#		if rnd==20:
#			plt.show()
		rnd=rnd+1
#		print "exl is:", exl
#		print "mi is:", mi
#		print "cf is:", cf
#		print "mxw is:", mxw
#		print "mrx is:", mrx
#		print "err is:", nerr
#		print "conv is:", conv
#		print "rnd is:", rnd
		yr=interp(w,arange(0,w.size)*2*pi/exl,real(y))
		yi=interp(w,arange(0,w.size)*2*pi/exl,imag(y))
	return yr+1j*yi

def spec2sigr(w,s,t):
	L=w.size
	N=t.size
	x=zeros(N)
	for i in range(0,L):
		x=x+real(s[i])*cos(w[i]*t)+imag(s[i])*sin(w[i]*t)
	return x	

def optimspec(t,x,w):
	sr=optimspecr(t,real(x),w)
	L=sr.size
	si=optimspecr(t,imag(x),w)
	if w[0]==0:
		ssr=zeros(L*2-1)*1j
		ssi=zeros(L*2-1)*1j
		ssr[L-1:2*L-1]=sr/2
		ssr[L-1]=sr[0]
		ssr[0:L-1]=sr[:0:-1].conj()/2
		ssi[L-1:2*L-1]=si/2
		ssi[L-1]=si[0]
		ssi[0:L-1]=si[:0:-1].conj()/2
#		print "real spec:", ssr
#		print "imag spec:", ssi
#		plt.plot(abs(ssr))
#		plt.plot(abs(ssi))
#		plt.show()
	else:
		ssr=zeros(L*2)*1j
		ssi=zeros(L*2)*1j
		ssr[L:2*L]=sr/2
		ssr[0:L]=sr[::-1].conj()/2
		ssi[L:2*L]=si/2
		ssi[0:L]=si[::-1].conj()/2
	spec=ssr+1j*ssi
#	plt.plot(abs(spec))
#	plt.show()
	return spec

def optimspec2(t,x,w):
	sr=optimspecr2(t,real(x),w)
	L=sr.size
	si=optimspecr2(t,imag(x),w)
	if w[0]==0:
		ssr=zeros(L*2-1)*1j
		ssi=zeros(L*2-1)*1j
		ssr[L-1:2*L-1]=sr/2
		ssr[L-1]=sr[0]
		ssr[0:L-1]=sr[:0:-1].conj()/2
		ssi[L-1:2*L-1]=si/2
		ssi[L-1]=si[0]
		ssi[0:L-1]=si[:0:-1].conj()/2
#		print "real spec:", ssr
#		print "imag spec:", ssi
#		plt.plot(abs(ssr))
#		plt.plot(abs(ssi))
#		plt.show()
	else:
		ssr=zeros(L*2)*1j
		ssi=zeros(L*2)*1j
		ssr[L:2*L]=sr/2
		ssr[0:L]=sr[::-1].conj()/2
		ssi[L:2*L]=si/2
		ssi[0:L]=si[::-1].conj()/2
	spec=ssr+1j*ssi
#	plt.plot(abs(spec))
#	plt.show()
	return spec

def optimspec4(t,x,w):
	sr=optimspecr4(t,real(x),w)
	L=sr.size
	si=optimspecr4(t,imag(x),w)
	if w[0]==0:
		ssr=zeros(L*2-1)*1j
		ssi=zeros(L*2-1)*1j
		ssr[L-1:2*L-1]=sr/2
		ssr[L-1]=sr[0]
		ssr[0:L-1]=sr[:0:-1].conj()/2
		ssi[L-1:2*L-1]=si/2
		ssi[L-1]=si[0]
		ssi[0:L-1]=si[:0:-1].conj()/2
#		print "real spec:", ssr
#		print "imag spec:", ssi
#		plt.plot(abs(ssr))
#		plt.plot(abs(ssi))
#		plt.show()
	else:
		ssr=zeros(L*2)*1j
		ssi=zeros(L*2)*1j
		ssr[L:2*L]=sr/2
		ssr[0:L]=sr[::-1].conj()/2
		ssi[L:2*L]=si/2
		ssi[0:L]=si[::-1].conj()/2
	spec=ssr+1j*ssi
#	plt.plot(abs(spec))
#	plt.show()
	return spec

def spec2sig(w,s,t):
	L=w.size
	N=t.size
	x=zeros(N)
	for i in range(0,L):
		x=x+s[i]*exp(1j*w[i]*t)
	return x

class fsinterpolator:
	def __init__(self,omega,coef):
		self.omega=omega
		self.coef=coef
	def __call__(self,x):
		L=self.omega.size
		N=x.size
		M=zeros((N,2*L))
		for i in range(0,L):
			M[:,i]=sin(self.omega[i]*x).T
			M[:,i+L]=cos(self.omega[i]*x).T
		return dot(M,self.coef.T)	  
def lsqspecr(t,x,w):
	if t.size!=x.size:
		print "The sizes of t and x don't match!"
		return False
	N=t.size
	L=w.size
	M=zeros((N,2*L))
        cf=zeros(2*L)
	for i in range(0,L):
		M[:,i]=sin(w[i]*t).T
		M[:,i+L]=cos(w[i]*t).T	
	if w[0]==0:	
		cf[1:]=linalg.lstsq(M[:,1:],x.T,rcond=1e-6)[0]
#		cf[1:]=LSQR(M[:,1:],x.T)[0]
		cf[0]=0
	else:
		cf=linalg.lstsq(M,x.T,rcond=1e-6)[0]
#		cf=LSQR(M,x.T)[0]
	return cf
		
def fsinterp(t,x,w):
	cf=lsqspecr(t,x,w)
	return fsinterpolator(w,cf)

if __name__=="__main__":
	t=arange(-79,80)
#	x=cos(pi/320*t)
	iq=genfromtxt("fake_gath.txt")
	iqf=fft.rfft(iq,axis=0)
	p=arange(0,2,4.0/t.size)
#	rp=genfromtxt("inputat1rowr.txt")
#	ip=genfromtxt("inputat1rowi.txt")
	x=iqf[20,:]
#	y=iqf[:,79]
#	plt.plot(abs(y))
#	plt.show()
#	plt.plot(abs(x))
#	plt.show()
#	x=rp+1j*ip
#	x=zeros(t.shape)
#	x[59]=1
#	plt.plot(t,real(x))
#	plt.plot(t,imag(x))
#	plt.show()
	w=p*pi*20/255
#	print w
	s=optimspec(t,real(x.T),w)
#	print s
	fw=zeros(2*w.size-1)
	fw[0:w.size-1]=-w[:0:-1]
	fw[w.size-1:]=w
	plt.plot(fw,abs(s))
	plt.show()
	
